package com.wave.wavesystem.ai.config;

import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.jdbc.DataSourceBuilder;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.core.JdbcTemplate;

import javax.sql.DataSource;

import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgDistanceType.COSINE_DISTANCE;
import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgIndexType.HNSW;

/**
 * 配置向量数据库 连接信息
 */
@Configuration
public class VectorStoreConfig {

    @Value("${wave.pg.url}")
    private String pgUrl;
    @Value("${wave.pg.username}")
    private String pgUsername;
    @Value("${wave.pg.password}")
    private String pgPassword;

    @Bean
    public VectorStore wavePgVectorStore(EmbeddingModel embeddingModel) {
        DataSource pgDataSource = DataSourceBuilder.create()
                .url(pgUrl)
                .username(pgUsername)
                .password(pgPassword)
                .build();
        JdbcTemplate pgJdbc = new JdbcTemplate(pgDataSource);

        return PgVectorStore.builder(pgJdbc, embeddingModel)
                .dimensions(1536)                    // Optional: defaults to model dimensions or 1536
                .distanceType(COSINE_DISTANCE)       // Optional: defaults to COSINE_DISTANCE
                .indexType(HNSW)                     // Optional: defaults to HNSW
                .initializeSchema(true)              // Optional: defaults to false
                .schemaName("wave_vector")                // Optional: defaults to "public"
                .vectorTableName("vector_store")     // Optional: defaults to "vector_store"
                .maxDocumentBatchSize(10000)         // Optional: defaults to 10000
                .build();

    }

}
